from easy_io import read_pkl_file, H5Writer, write_pkl_file
import numpy as np
import h5py
from scipy.ndimage import zoom


def read_partly_from_h5_file(h5_file_obj):
    def helper(path_starts_shape):
        path, starts, shape = path_starts_shape
        slices = tuple(slice(max(st, 0), st + sp) for st, sp in zip(starts, shape))
        raw = h5_file_obj[path][slices]
        mshape = h5_file_obj[path].shape
        pad_width = np.asarray([(0 - st, st + sp - msp) for st, sp, msp in zip(starts, shape, mshape)], 'int')
        pad_width[pad_width < 0] = 0
        paded = np.pad(raw, pad_width, 'constant', constant_values=-1024)
        assert paded.shape == tuple(shape)
        return paded

    return helper


def gen(candidate_pkl_file, data_h5_file, min_edge_length):
    candidates = read_pkl_file(candidate_pkl_file)

    with h5py.File(data_h5_file, 'r') as f:
        reader = read_partly_from_h5_file(f)
        for i, c in enumerate(candidates):
            c['index'] = i
            c['path'] = str(i)

            bbox = np.asarray(c['bbox'], 'int')
            shape = bbox[-3:] - bbox[:3]
            center = (bbox[-3:] + bbox[:3] - 1) / 2
            diameter = np.max(shape[:2])
            c['center'] = center
            c['diameter'] = diameter

            scanid = c['scanid']
            spacing = np.asarray(f[scanid]['vol'].attrs['spacing'])
            assert spacing[0] == spacing[1], scanid

            unit = np.min(spacing)
            edge_length = max(min_edge_length, diameter)
            crop_shape = np.asarray(np.ceil((edge_length - 1) * unit / spacing + 1), 'int')
            start = np.asarray(np.floor(center - crop_shape / 2 + 1), 'int')
            vol = reader((scanid+'/vol', start, crop_shape))

            vol = zoom(vol, min_edge_length / np.asarray(vol.shape), output='float32', order=1, mode='nearest')
            assert np.all(np.asarray(vol.shape) == min_edge_length)
            yield str(i), vol

    write_pkl_file(candidate_pkl_file, candidates)


if __name__ == '__main__':
    # H5Writer(
    #     '/ssd_1t/wangd/kaggle_data/LIDC_true_nodules_vol.hdf5',
    #     'w',
    #     gen(
    #         candidate_pkl_file='/ssd_1t/wangd/kaggle_data/LIDC_true_nodules.pkl',
    #         data_h5_file='/data_4t/Kaggle/backup/lidc/vol.hdf5',
    #         min_edge_length=68,
    #     )
    # )
    # H5Writer(
    #     '/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v1.hdf5',
    #     'w',
    #     gen(
    #         candidate_pkl_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v1.pkl',
    #         data_h5_file='/data_4t/Kaggle/lidc&kaggle/train_data.hdf5',
    #         min_edge_length=68,
    #     )
    # )
    # H5Writer(
    #     '/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v2.hdf5',
    #     'w',
    #     gen(
    #         candidate_pkl_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v2.pkl',
    #         data_h5_file='/data_4t/Kaggle/lidc&kaggle/train_data.hdf5',
    #         min_edge_length=68,
    #     )
    # )
    H5Writer(
        '../../data/reduced_candidates/LIDC_true_nodules_vol.hdf5',
        'w',
        gen(
            candidate_pkl_file='../../data/reduced_candidates/LIDC_true_nodules.pkl',
            data_h5_file='../../data/lidc/vol.hdf5',
            min_edge_length=68,
        )
    )
